# FaissKNeighbors.py

import faiss
import numpy as np

class FaissKNeighbors:
    def __init__(self, k=5, gpu=False):
        self.k = k
        self.gpu = gpu
        self.index = None
        self.res = None

    def fit(self, X, y):
        self.y = y
        dim = X.shape[1]
        self.index = faiss.IndexFlatL2(dim)
        
        if self.gpu:
            print("正在初始化FAISS的GPU资源...")
            self.res = faiss.StandardGpuResources()
            self.index = faiss.index_cpu_to_gpu(self.res, 0, self.index)
            print("FAISS的GPU资源初始化完成。")
        
        self.index.add(X)

    def predict(self, X):
        distances, indices = self.index.search(X, self.k)
        votes = self.y[indices]
        predictions = np.array([np.argmax(np.bincount(vote)) for vote in votes])
        return predictions

    def score(self, X, y):
        predictions = self.predict(X)
        accuracy = np.mean(predictions == y)
        return accuracy